import random
import torch
import numpy as np
import gym
from gym import spaces
from skidl.schematics.geometry import BBox, Point, Tx, Vector
import datetime
Array = np.ndarray
from .schDraw import calc_dist_of_net
class Part(gym.Env):
    # This method is called when an object is created.
    # It's purpose is to initialize the object.
    def __init__(self,
                 part,
                 max_step,
                 bboxHeight,
                 bboxWidth,
                 env):

        self.part = part

        self.max_steps = max_step
        self.steps_done = 0

        self.bboxHeight = int(bboxHeight)
        self.bboxWidth = int(bboxWidth)
        self.center_x = self.bboxWidth / 2
        self.center_y = self.bboxHeight / 2
        self.env = env

        # reset
        self.pos_x = None
        self.pos_y = None
        self.overlap_force_x = None
        self.overlap_force_y = None

        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        '''action and obs'''
        # obs_space = {
        #     "overlap_force_x": spaces.Box(low=-self.bboxWidth, high=self.bboxWidth, shape=(1,),
        #                       dtype=np.float32), #避免重叠驱动力，x方向
        #     "overlap_force_y": spaces.Box(low=-self.bboxHeight, high=self.bboxHeight, shape=(1,),
        #                      dtype=np.float32), #避免重叠驱动力，y方向
        #     "position_x": spaces.Box(low=0.0, high=self.bboxWidth, shape=(1,),
        #                            dtype=np.float32), #x坐标
        #     "position_y": spaces.Box(low=0.0, high=self.bboxHeight, shape=(1,),
        #                              dtype=np.float32), #y坐标
        # }
        # self.observation_space = spaces.Dict(obs_space)
        self.observation_space = spaces.Box(
            low=np.array([-np.inf, np.inf, 0.0, 0.0],dtype=np.float32),
            high=np.array([-np.inf, np.inf, self.bboxWidth, self.bboxHeight],dtype=np.float32),
        )
        self.action_space = spaces.Box(
            low=np.array([-10, -10], dtype=np.float32),
            high=np.array([10, 10],dtype=np.float32))

    def reset(self):
        self.steps_done = 0

        # 将元件位置随机放置
        pt = Point(random.random() * self.bboxWidth, random.random() * self.bboxHeight)
        self.part.tx = self.part.tx.move(pt)

        return self.get_state()

    def step(self,action: Array) :
        self.steps_done += 1
        # 加上边界限制
        # if((self.part.place_bbox * self.part.tx).ctr.x+action[0] > self.bboxWidth):
        #     action[0] = 800 - (self.part.place_bbox * self.part.tx).ctr.x
        # elif((self.part.place_bbox * self.part.tx).ctr.x+action[0] < 0):
        #     action[0] = 800-(self.part.place_bbox * self.part.tx).ctr.x
        #
        # if ((self.part.place_bbox * self.part.tx).ctr.y + action[1] > self.bboxHeight):
        #     action[1] = 800 - (self.part.place_bbox * self.part.tx).ctr.y
        # elif ((self.part.place_bbox * self.part.tx).ctr.y + action[1] < 0):
        #     action[1] = 800-(self.part.place_bbox * self.part.tx).ctr.y
        pt_temp = Point(action[0], action[1])

        self.part.tx = self.part.tx.move(pt_temp)

        next_state = self.get_state()
        reward = self.get_reward()  # reward的计算要在get_state之后，确保计算出来的reward是最新的
        done = torch.zero(dtype=torch.bool).to(self.device)
        info = None
        return next_state, reward, done, info

    def get_state(self) :
        self.pos_x = (self.part.place_bbox * self.part.tx).ctr.x
        self.pos_y = (self.part.place_bbox * self.part.tx).ctr.y

        overlap_force_vec = overlap_force(self.part, self.env.parts)
        self.overlap_force_x = overlap_force_vec.x
        self.overlap_force_y = overlap_force_vec.y
        npArr =  np.hstack((self.overlap_force_x, self.overlap_force_y, self.pos_x, self.pos_y))
        return torch.from_numpy(npArr).float().to(self.device)

    def get_reward(self):
        dist_2_centX = abs(self.center_x - self.pos_x)
        dist_2_centY = abs(self.center_y - self.pos_y)
        abs_overlap_x = abs(self.overlap_force_x)
        abs_overlap_y = abs(self.overlap_force_y)
        n = -1
        m = -4
        reward = n * dist_2_centX + n * dist_2_centY + m * abs_overlap_x + m * abs_overlap_y

        return torch.Tensor(reward)


def overlap_force(part, parts):
    """Compute the repulsive force on a part from overlapping other parts.

    Args:
        part (Part): Part affected by forces from other overlapping parts.
        parts (list): List of parts to check for overlaps.

    Returns:
        Vector: Force upon given part.
    """

    # Bounding box of given part.
    part_bbox = part.place_bbox * part.tx

    # Compute the overlap force of the bbox of this part with every other part.
    total_force = Vector(0, 0)
    for other_part in set(parts) - {part}:
        other_part_bbox = other_part.place_bbox * other_part.tx

        # No force unless parts overlap.
        if part_bbox.intersects(other_part_bbox):
            # Compute the movement needed to separate the bboxes in left/right/up/down directions.
            # Add some small random offset to break symmetry when parts exactly overlay each other.
            # Move right edge of part to the left of other part's left edge, etc...
            moves = []
            rnd = Vector(random.random() - 0.5, random.random() - 0.5)
            for edges, dir in ((("ll", "lr"), Vector(1, 0)), (("ul", "ll"), Vector(0, 1))):
                move = (getattr(other_part_bbox, edges[0]) - getattr(part_bbox, edges[1]) - rnd) * dir
                moves.append([move.magnitude, move])
                # Flip edges...
                move = (getattr(other_part_bbox, edges[1]) - getattr(part_bbox, edges[0]) - rnd) * dir
                moves.append([move.magnitude, move])

            # Select the smallest move that separates the parts.
            move = min(moves, key=lambda m: m[0])

            # Add the move to the total force on the part.
            total_force += move[1]

    return total_force
